package leetcode;

import java.awt.*;
import java.util.Arrays;

/**
 * 最小生成树
 */
public class NC159MinSpaningTree {

    /**
     * Prim算法
     * @param n
     * @param m
     * @param cost
     * @return
     */
    public int miniSpanningTree (int n, int m, int[][] cost) {
        int minCost = 0; // 生成树最小代价
        boolean[] inTree = new boolean[n + 1]; //记录节点i是否在生成树中(索引从1开始)
        int inTreeCount = 0; //记录在生成树中的节点个数
        int[] dist = new int[n + 1]; // dist[i]记录节点i到生成树集合的最小代价(索引从1开始)
        int[][] graph = new int[n + 1][n + 1];
        for (int i = 1; i <= n; i++) {
            for (int j = 1; j <= n; j++) {
                graph[i][j] = Integer.MAX_VALUE;
            }
        }

        for (int i = 0; i < m; i++) { // 测试案例中存在同一个边，权重不一样的情况，取较小的权重值
            if (cost[i][2] < graph[cost[i][0]][cost[i][1]]) {
                graph[cost[i][0]][cost[i][1]] = cost[i][2];
            }
            if (cost[i][2] < graph[cost[i][1]][cost[i][0]]) {
                graph[cost[i][1]][cost[i][0]] = cost[i][2];
            }
        }

        //初始选节点1到生成树
        inTree[1] = true;
        inTreeCount = 1;
        for (int i = 1; i <= n; i++) { //更新其它节点到生成树的距离（即节点1）
            dist[i] = graph[1][i];
        }
        while (inTreeCount < n) {
            int minDist = Integer.MAX_VALUE;
            int index = 0;
            for (int i = 1; i <= n; i++) {
                if (!inTree[i] && dist[i] < minDist) {
                    minDist = dist[i];
                    index = i;
                }
            }
            inTree[index] = true;
            inTreeCount ++;
            minCost += dist[index];

            for (int i = 1; i <= n; i++) {
                if (!inTree[i] && graph[index][i] < dist[i]) {
                    dist[i] = graph[index][i];
                }
            }
        }
        return minCost;
    }

    //kruskal算法
    public int miniSpanningTree2 (int n, int m, int[][] cost) {

        int minCost = 0;
        Arrays.sort(cost, (a, b) -> a[2] - b[2]);
        UnionFindSet unionFindSet = new UnionFindSet(n + 1);
        for (int i = 0; i < m; i++) {
            if (!unionFindSet.isConnected(cost[i][0],cost[i][1])) {
                unionFindSet.union(cost[i][0],cost[i][1]);
                minCost += cost[i][2];
            }
        }
        return minCost;
    }

    class UnionFindSet {
        int[] parent;
        int count;

        public UnionFindSet(int n) {
            this.parent = new int[n];
            this.count = n ;
            for (int i = 0; i < n; i++) {
                this.parent[i] = i;
            }
        }

        public int find(int x) {
            while (x != parent[x])
                x = parent[x];
            return x;
        }

        public void union(int x, int y) {
            int xParen = find(x);
            int yParent = find(y);
            if (xParen != yParent) {
                this.parent[xParen] = yParent;
                this.count --;
            }
        }

        public boolean isConnected(int x, int y) {
            return find(x) == find(y);
        }
    }


}
